{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from pprint import pprint\n",
    "import threading\n",
    "\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain.globals import set_llm_cache\n",
    "from langchain_core.caches import InMemoryCache\n",
    "import pandas as pd\n",
    "import promptquality as pq\n",
    "from tqdm.auto import tqdm\n",
    "from llm_handler import LLMHandler\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "set_llm_cache(InMemoryCache())\n",
    "\n",
    "load_dotenv(\"../.env\")\n",
    "pq.login(\"console.demo.rungalileo.io\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_handler = LLMHandler()\n",
    "model = \"command-a-03-2025\"\n",
    "llm = llm_handler.get_llm(model, temperature=0.0, max_tokens=4000)\n",
    "# llm.invoke(\"Hello, how are you?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = glob(\"../data/datasets/*.parquet\")\n",
    "print(len(files))\n",
    "files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_name = \"agent-lb-v1\"\n",
    "system_msg = {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": 'Your job is to use the given tools to answer the query of human. If there is no relevant tool then reply with \"I cannot answer the question with given tools\". If tool is available but sufficient information is not available, then ask human to get the same. You can call as many tools as you want. Use multiple tools if needed. If the tools need to be called in a sequence then just call the first tool.',\n",
    "        }\n",
    "\n",
    "chainpoll_tool_selection_scorer = pq.CustomizedChainPollScorer(\n",
    "                scorer_name=pq.CustomizedScorerName.tool_selection_quality,\n",
    "                model_alias=pq.Models.gpt_4o,\n",
    "            )\n",
    "\n",
    "project_id = pq.get_project_from_name(project_name)\n",
    "\n",
    "def check_run_exists(model_name, file_name):\n",
    "    run_name = f\"{model_name} {file_name}\"    \n",
    "    try: \n",
    "        pq.get_run_from_name(run_name, project_id.id)\n",
    "        print(f\"Run {run_name} already exists\")\n",
    "        return True\n",
    "    except ValueError:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# models = llm_handler.available_models[\"together\"]\n",
    "# models = llm_handler.get_all_models()\n",
    "models = [\n",
    "            # \"nvidia/llama-3.1-nemotron-nano-8b-v1\",\n",
    "            # \"nvidia/llama-3.3-nemotron-super-49b-v1\"\n",
    "            ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_lock = threading.Lock()\n",
    "\n",
    "def process_file(model_name, file_path):\n",
    "    file_name = Path(file_path).stem\n",
    "    run_name = f\"{model_name} {file_name}\"\n",
    "    \n",
    "    # Check if run already exists\n",
    "    try:\n",
    "        pq.get_run_from_name(run_name, project_id.id)\n",
    "        return model_name, file_name, None\n",
    "    except ValueError:\n",
    "        pass  # Run doesn't exist, continue processing\n",
    "    \n",
    "    df = pd.read_parquet(file_path, engine=\"fastparquet\")\n",
    "\n",
    "    evaluate_handler = pq.GalileoPromptCallback(\n",
    "        project_name=project_name,\n",
    "        run_name=run_name,\n",
    "        scorers=[chainpoll_tool_selection_scorer],\n",
    "    )\n",
    "\n",
    "    outputs = []\n",
    "    llm = llm_handler.get_llm(model_name, max_tokens=4000)\n",
    "    for row in tqdm(df.itertuples(), \n",
    "                   desc=run_name, \n",
    "                   total=len(df), \n",
    "                   position=None, \n",
    "                   leave=False):\n",
    "        try:\n",
    "            chain = llm.bind_tools(row.tools_langchain).with_retry(\n",
    "                wait_exponential_jitter=True,\n",
    "                stop_after_attempt=4,\n",
    "            )\n",
    "            output = chain.invoke(\n",
    "                [system_msg, *row.conversation], \n",
    "                config=dict(callbacks=[evaluate_handler])\n",
    "            )\n",
    "            outputs.append(output)\n",
    "        except Exception as e:\n",
    "            with print_lock:\n",
    "                print(f\"\\nSkipping sample in {run_name} due to error: {str(e)}\")\n",
    "            continue\n",
    "\n",
    "    evaluate_handler.finish()\n",
    "    return model_name, file_name, outputs\n",
    "\n",
    "def process_model(model_name):\n",
    "    with ThreadPoolExecutor(max_workers=3) as file_executor:\n",
    "        file_futures = [\n",
    "            file_executor.submit(process_file, model_name, file_path)\n",
    "            for file_path in files\n",
    "        ]\n",
    "        \n",
    "        results = []\n",
    "        for future in tqdm(file_futures, \n",
    "                          desc=f\"Files for {model_name}\", \n",
    "                          position=None, \n",
    "                          leave=False):\n",
    "            results.append(future.result())\n",
    "    return results\n",
    "\n",
    "# Main execution\n",
    "with ThreadPoolExecutor(max_workers=3) as model_executor:\n",
    "    # Submit each model to the executor\n",
    "    model_futures = [model_executor.submit(process_model, model) for model in models]\n",
    "    \n",
    "    # Collect all results\n",
    "    all_results = []\n",
    "    for future in tqdm(model_futures, \n",
    "                      desc=\"Overall Models Progress\", \n",
    "                      position=0, \n",
    "                      leave=True):\n",
    "        all_results.extend(future.result())\n",
    "\n",
    "# Print summary\n",
    "completed_runs = [(model, file) for model, file, outputs in all_results if outputs is not None]\n",
    "if completed_runs:\n",
    "    print(\"\\nNewly completed runs:\")\n",
    "    for model, file in completed_runs:\n",
    "        print(f\"- {file} with {model}\")\n",
    "else:\n",
    "    print(\"\\nNo new runs were processed - all runs already existed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langgraph",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
